#ifndef ASYNCWEBSOCKET_H_
#define ASYNCWEBSOCKET_H_

#include "AsyncTCP.hpp"
#include "ESPAsyncWebServer.h"
#include "AsyncWebSynchronization.h"

#define CONFIG_WS_MAX_QUEUE_MESSAGES    32
#define CONFIG_MAX_WS_CLIENTS   8

class AsyncWebSocket;
class AsyncWebSocketResponse;
class AsyncWebSocketClient;
class AsyncWebSocketControl;

using AwsFrameInfo = struct {
    uint8_t     message_opcode;     // 消息操作码，标识整个消息的类型（WS_TEXT/WS_BINARY）
    uint32_t    num;                // 消息分片时，当前帧在消息中的编号
    uint8_t     final;              // 消息的最后一帧标志（1：最后一帧）
    uint8_t     masked;             // 当前帧是否被掩码保护
    uint8_t     opcode;             // 当前帧操作码
    uint64_t    len;                // 当前帧的有效载荷长度
    uint8_t     mask[4];            // 帧被掩码保护时，掩码密钥
    uint64_t    index;              // 当前帧有效载荷数据量的偏移量
};

typedef enum {
    WS_DISCONNECTED,    // 客户端未连接
    WS_CONNECTED,       // 客户端已成功连接
    WS_DISCONNECTING    // 客户端正在断开连接
} AwsClientStatus;
typedef enum {
    WS_CONTINUATION,        // 继续帧
    WS_TEXT,                // 文本帧
    WS_BINARY,              // 二进制帧
    WS_DISCONNECT = 0x08,   // 关闭帧
    WS_PING,                // Ping帧
    WS_PONG                 // Pong帧
} AwsFrameType;
typedef enum {
    WS_MSG_SENDING,     // 消息分片发送中
    WS_MSG_SENT,        // 帧间状态
    WS_MSG_ERROR        // 发送过程发生错误
} AwsMessageStatus;
typedef enum {
    WS_EVT_CONNECT,     // 客户端成功连接
    WS_EVT_DISCONNECT,  // 连接被断开
    WS_EVT_PONG,        // 收到服务器的Pong响应
    WS_EVT_ERROR,       // 发生错误
    WS_EVT_DATA         // 接收至数据（文本/二进制）
} AwsEventType;

// 用于接收、发送时数据存放的类
class AsyncWebSocketMessageBuffer {
private:
    uint8_t*    m_data;     // 实际存储数据的内存指针
    size_t      m_len;      // 存储的数据长度
    bool        m_lock;     // 锁标志，为true时其他操作会被阻塞
    uint32_t    m_count;    // 引用计数器，记录引用当前引用该缓冲区的对象数量
public:
    AsyncWebSocketMessageBuffer();
    AsyncWebSocketMessageBuffer(size_t size);
    AsyncWebSocketMessageBuffer(uint8_t* data, size_t size);
    AsyncWebSocketMessageBuffer(const AsyncWebSocketMessageBuffer &);
    AsyncWebSocketMessageBuffer(AsyncWebSocketMessageBuffer &&);
    ~AsyncWebSocketMessageBuffer();
    void operator ++(int i)
    {
        m_count++;
    }
    void operator --(int i)
    {
        if (m_count > 0) {
            m_count--;
        }
    }
    bool reserve(size_t size);
    void lock()
    {
        m_lock = true;
    }
    void unlock()
    {
        m_lock = false;
    }
    uint8_t* get()
    {
        return m_data;
    }
    size_t length()
    {
        return m_len;
    }
    uint32_t count()
    {
        return m_count;
    }
    bool canDelete()
    {
        return (!m_count && !m_lock);
    }

    friend AsyncWebSocket;
};

// WebSocket消息的基类
class AsyncWebSocketMessage {
protected:
    uint8_t             m_opcode;   // WebSocket帧操作码
    bool                m_mask;     // 消息是否要掩码处理
    AwsMessageStatus    m_status;   // 消息的发送状态
public:
    AsyncWebSocketMessage()
        : m_opcode(WS_TEXT)
        , m_mask(false)
        , m_status(WS_MSG_ERROR)
    {}
    virtual ~AsyncWebSocketMessage() {}
    virtual void ack(size_t len, uint64_t time) {}
    virtual size_t send(AsyncClient* client)
    {
        return 0;
    }
    virtual bool finished()
    {
        return m_status != WS_MSG_SENDING;
    }
    virtual bool betweenFrames() const
    {
        return false;
    }
};

// 基本WebSocket消息类
class AsyncWebSocketBasicMessage : public AsyncWebSocketMessage {
private:
    size_t  m_len;      // 消息总长度
    size_t  m_sent;     // 已成功发送的字节数（不含头部和掩码）
    size_t  m_ack;      // 需要确认的字节数
    size_t  m_acked;    // 已收到确认的字节灵长
    uint8_t*    m_data; // 消息指针
public:
    AsyncWebSocketBasicMessage(const char* data, size_t len, uint8_t opcode = WS_TEXT, bool mask = false);
    AsyncWebSocketBasicMessage(uint8_t opcode = WS_TEXT, bool mask = false);
    virtual ~AsyncWebSocketBasicMessage() override;
    virtual bool betweenFrames() const override
    {
        return m_acked == m_ack;
    }
    virtual void ack(size_t len, uint64_t time) override;
    virtual size_t send(AsyncClient* client) override;
};

// 多播WebSocket消息类
class AsyncWebSocketMultiMessage : public AsyncWebSocketMessage {
private:
    size_t  m_len;      // 消息总长度
    size_t  m_sent;     // 已成功发送的字节数（不含头部和掩码）
    size_t  m_ack;      // 需要确认的字节数
    size_t  m_acked;    // 已收到确认的字节灵长
    uint8_t*    m_data; // 消息指针
    AsyncWebSocketMessageBuffer*    m_WSbuffer;
public:
    AsyncWebSocketMultiMessage(AsyncWebSocketMessageBuffer* buffer, uint8_t opcode = WS_TEXT, bool mask = false);
    virtual ~AsyncWebSocketMultiMessage() override;
    virtual bool betweenFrames() const override
    {
        return m_acked == m_ack;
    }
    virtual void ack(size_t len, uint64_t time) override;
    virtual size_t send(AsyncClient* client) override;
};

using AwsEventHandler = std::function<void(AsyncWebSocket*, AsyncWebSocketClient*, AwsEventType, void* arg, uint8_t* data, size_t len)>;
class AsyncWebSocket : public AsyncWebHandler {
public:
    using AsyncWebSocketClientLinkedList = LinkedList<AsyncWebSocketClient*>;

    AsyncWebSocket(const std::string &url);
    ~AsyncWebSocket();
    // 获取WebSocket绑定的URL
    const char* url() const
    {
        return m_url.c_str();
    }
    // 设置WebSocket功能是否启用
    void enable(bool e)
    {
        m_enabled = e;
    }
    // 检查WebSocket功能启用状态
    bool enabled() const
    {
        return m_enabled;
    }
    bool availableForWriteAll();
    bool availableForWrite(uint32_t id);

    size_t count() const;
    AsyncWebSocketClient* client(uint32_t id);
    bool hasClient(uint32_t id)
    {
        return client(id) != nullptr;
    }

    void close(uint32_t id, uint16_t code = 0, const char* message = nullptr);
    void closeAll(uint16_t code = 0, const char* message = nullptr);
    void cleanupClients(uint16_t maxClients = CONFIG_MAX_WS_CLIENTS);

    void ping(uint32_t id, uint8_t *data = nullptr, size_t len = 0);
    void pingAll(uint8_t *data = nullptr, size_t len = 0);

    void text(uint32_t id, const char* message);
    void text(uint32_t id, const std::string &message);

    void textAll(const char* message);
    void textAll(AsyncWebSocketMessageBuffer* buffer);

    void binary(uint32_t id, const char* message);
    void binary(uint32_t id, const std::string &message);

    void binaryAll(const char* message);
    void binaryAll(AsyncWebSocketMessageBuffer* buffer);

    void message(uint32_t id, AsyncWebSocketMessage* message);
    void messageAll(AsyncWebSocketMessage* message);

    // 设置WebSkocet事件处理函数
    void onEvent(AwsEventHandler h)
    {
        m_eventHandler = h;
    }

    uint32_t _getNextId()
    {
        return m_cNextID++;
    }
    void _addClient(AsyncWebSocketClient* client);
    void _handleDisconnect(AsyncWebSocketClient* client);
    void _handleEvent(AsyncWebSocketClient* client, AwsEventType type, void* arg, uint8_t* data, size_t len);
    virtual bool canHandle(AsyncWebServerRequest* req) override final;
    virtual void handleRequest(AsyncWebServerRequest* req) override final;

    AsyncWebSocketMessageBuffer* makeBuffer(size_t size = 0);
    AsyncWebSocketMessageBuffer* makeBuffer(uint8_t* data, size_t size);
    LinkedList<AsyncWebSocketMessageBuffer*> m_buffers;
    void _cleanBuffers();

    AsyncWebSocketClientLinkedList getClients() const;

private:
    std::string     m_url;          // WebSocket服务绑定的URL路径
    AsyncWebSocketClientLinkedList  m_clients;  // 连接WebSocket服务的客户端对象
    uint32_t        m_cNextID;      // ID计数器（为每个新连接的客户端分配唯一标识符）
    AwsEventHandler m_eventHandler; // 用户注册的事件回调（处理客户端连接/断开/消息事件）
    bool            m_enabled;      // 标识WebSocket服务是否处于启用状态
    AsyncWebLock    m_lock;         // 保护共享资源的线程锁

};

class AsyncWebSocketClient {
private:
    AsyncClient*    m_client;   // 当前客户端的网络连接
    AsyncWebSocket* m_server;   // 当前客户端连接的服务器
    uint32_t        m_clientID; // 客户端连接的服务器
    AwsClientStatus m_status;   // 客户端当前状态

    LinkedList<AsyncWebSocketControl*> m_controlQueue;  // 控制帧队列
    LinkedList<AsyncWebSocketMessage*> m_messageQueue;  // 消息队列

    uint8_t         m_pstate;   // 处理当前帧所处的阶段（初始状态/解析数据）
    AwsFrameInfo    m_pinfo;    // 当前帧的元信息（操作码、帧长度、是否分片、Mask等）

    uint64_t    m_lastMessageTime;  // 记录最后一次成功发送/接收消息时间戳
    uint64_t    m_keepAlivePeriod;  // 心跳时间间隔(ms)

    void m_queueMessage(AsyncWebSocketMessage* dataMessage);    //
    void m_queueControl(AsyncWebSocketControl* controlMessage); //
    void m_runQueue();                                          //
public:
    void*   _tempObject;

    AsyncWebSocketClient(AsyncWebServerRequest* req, AsyncWebSocket* server);
    ~AsyncWebSocketClient();

    uint32_t id()
    {
        return m_clientID;
    }
    AwsClientStatus status()
    {
        return m_status;
    }
    AsyncClient* client()
    {
        return m_client;
    }
    AsyncWebSocket* server()
    {
        return m_server;
    }
    AwsFrameInfo const &pinfo() const
    {
        return m_pinfo;
    }

    ip_addr_t remoteIP();
    uint16_t remotePort();

    void close(uint16_t code = 0, const char* message = nullptr);
    void ping(uint8_t* data = nullptr, size_t len = 0);

    // 设置keepAlive时间间隔
    void keepAlivePeriod(uint16_t seconds)
    {
        m_keepAlivePeriod = seconds * 1000;
    }
    // 获取keepAlived时间间隔
    uint16_t keepAlivePeriod()
    {
        return (uint16_t)(m_keepAlivePeriod / 1000);
    }

    // 将消息添加至队列并尝试发送
    void message(AsyncWebSocketMessage* message)
    {
        m_queueMessage(message);
    }
    bool queueIsFull();

    void text(const char* message);
    void text(uint8_t* message, size_t len);
    void text(AsyncWebSocketMessageBuffer* buffer);

    void binary(const char* message);
    void binary(uint8_t* message, size_t len);
    void binary(AsyncWebSocketMessageBuffer* buffer);

    // 消息队列不满则返回true
    bool canSend()
    {
        return m_messageQueue.length() <  CONFIG_WS_MAX_QUEUE_MESSAGES;
    }

    void _onAck(size_t len, uint64_t time);
    void _onError(err_t) {}
    void _onPoll();
    void _onTimeout(uint64_t time);
    void _onDisconnect();
    void _onData(void* pbuf, size_t plen);
};

class AsyncWebSocketResponse : public AsyncWebServerResponse {
public:
    AsyncWebSocketResponse(const std::string &key, AsyncWebSocket* server);
    void _respond(AsyncWebServerRequest* req);
    size_t _ack(AsyncWebServerRequest* req, size_t len, uint64_t time);
    bool _sourceValid() const
    {
        return true;
    }
private:
    std::string     m_content;
    AsyncWebSocket* m_server;
};

#endif
